package io.xccit.deepseek.service.impl;

import io.xccit.deepseek.entity.ChatEntity;
import io.xccit.deepseek.enums.MessageType;
import io.xccit.deepseek.service.OllamaService;
import io.xccit.deepseek.utils.SseServer;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.ChatResponse;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.ollama.OllamaChatClient;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.util.Map;
import java.util.stream.Collectors;

/**
 * @author CH_ywx
 * @version 1.0
 * @create 2025/4/7
 * <p>OllamaServiceImpl</p>
 */
@Slf4j
@Service
public class OllamaServiceImpl implements OllamaService {

    private final OllamaChatClient chatClient;

    @Autowired
    public OllamaServiceImpl(OllamaChatClient chatClient) {
        this.chatClient = chatClient;
    }


    /**
     * 聊天
     *
     * @param message 消息
     * @return 聊天结果
     */
    @Override
    public Map<String, Object> chat(String message) {
        ChatResponse chatResponse = chatClient.call(new Prompt(new UserMessage(message)));
        String response = chatResponse.getResult().getOutput().getContent();
        return Map.of("response", response);
    }

    /**
     * 聊天流式返回
     *
     * @param message 消息
     * @return 聊天结果
     */
    @Override
    public Flux<ChatResponse> chatStream(String message) {
        Prompt prompt = new Prompt(new UserMessage(message));
        Flux<ChatResponse> chatResponseFlux = chatClient.stream(prompt);
        return chatResponseFlux;
    }

    /**
     * 聊天医生
     *
     * @param chatEntity
     */
    @Override
    public void chatDoctor(ChatEntity chatEntity) {
        String userId = chatEntity.getCurrentUserName();
        String message = chatEntity.getMessage();
        Flux<ChatResponse> chatResponseFlux = chatClient.stream(new Prompt(new UserMessage(message)));
        chatResponseFlux.subscribe(chatResponse -> {
            String content = chatResponse.getResult().getOutput().getContent();
            SseServer.sendMessage(userId, content, MessageType.MESSAGE);
        });
        SseServer.sendMessage(userId, "finish", MessageType.FINISH);
    }
}
